import json
import logging
import numpy as np
from paramiko import SSHClient, AutoAddPolicy
from PIL import Image
import re
import requests
import random
import math


def find_closest(coord1, coord2, target_coord):
    """
    找出两个二维坐标中与第三个二维坐标最近的那个。

    参数:
    coord1: tuple，表示第一个二维坐标，例如 (x1, y1)
    coord2: tuple，表示第二个二维坐标，例如 (x2, y2)
    target_coord: tuple，表示目标二维坐标，例如 (xt, yt)

    返回:
    tuple，最近的坐标 (x, y)
    """
    def distance(c1, c2):
        """计算两个二维坐标之间的欧几里得距离"""
        return math.sqrt((c1[0] - c2[0])**2 + (c1[1] - c2[1])**2)
    
    # 计算与目标坐标的距离
    dist1 = distance(coord1, target_coord)
    dist2 = distance(coord2, target_coord)

    # 返回距离较小的坐标
    return coord1 if dist1 <= dist2 else coord2


class PostAPI:
    """ 调用远程的VLM API, 通过工作环境的初始视觉状态和指令获取用以完成指令的抓取位置。
        由于 VIMA 的图片信息通过列表形式表示，因此需要将列表表示的图片转换为图片并保存在本地。
        由于 VLM 的输入为图片的路径，且图片需要和 VLM 在同一个设备上，因此需要将图片上传至远程服务器的指定位置。
        Args:
            url: API 的 url
            remote_ip: 远程服务器的 ip
            remote_user_name: 远程服务器的用户名
            remote_password: 远程服务器的密码
            port: 远程服务器的端口
            remote_path: 远程服务器上保存图片的文件夹路径
            local_path: 本地保存图片的文件夹路径

        Example:
            >>>  post_api = PostAPI(
                    url = "http://127.0.0.1:8080/qwen2_vl_service/chat",
                    remote_ip="connect.bjc1.seetacloud.com",
                    remote_user_name="root",
                    port=40115,
                    remote_password="mfH6cr4mtaPe",
                    remote_path="/root/autodl-tmp/qwen2_vl_service/images",
                    local_path="/home/zp/文档/code/robovlm/local_images"
                )
            >>> init_image = [
                    [(255, 0, 0), (0, 255, 0), (0, 0, 255)],  # 第一行像素 (红, 绿, 蓝)
                    [(255, 255, 0), (255, 0, 255), (0, 255, 255)], # 第二行像素
                    [(128, 128, 128), (64, 64, 64), (0, 0, 0)]   # 第三行像素 (灰色, 深灰, 黑)
                ]
            >>> action = post_api.post_request(init_image=init_image, prompt="put the pen in the box")
    """
    def __init__(
            self,
            url: str = None,
            remote_ip: str = None,
            port: int = 22,
            remote_user_name: str = None,
            remote_password: str = None,          
            remote_path: str = None,
            local_path: str = None
    ):
        self.url = url
        self.remote_ip = remote_ip
        self.port = port
        self.remote_user_name = remote_user_name
        self.remote_password = remote_password
        self.remote_path = remote_path
        self.local_path = local_path
    
    def save_image_to_local(self,
                            init_image: list[list[int]],
                            prompt: str
                            ):
        
        """将列表类型表示的图片转换为 Image 类型并保存在本地.
        Args:
            init_image_list (list[list[int]]): 列表类型表示的图片.
            prompt (str): 用户发布给机器人的指令，用于图片的命名.
        """
        image = Image.fromarray(init_image)
        # 保存图片到本地，文件名使用指令来命名。
        self.image_save_path = self.local_path + f"/{prompt}.jpg"
        # 保存为 JPG 格式。
        image.save(self.image_save_path)

    def upload_image_to_remote(self, prompt: str):
        """使用 ssh 和 ftp 上传图片到服务器"""
        # 创建 SSH 客户端.
        logging.getLogger("paramiko").setLevel(logging.CRITICAL)
        ssh = SSHClient()
        ssh.set_missing_host_key_policy(AutoAddPolicy())
        try:
            # 连接到远程主机
            ssh.connect(self.remote_ip, self.port, self.remote_user_name, self.remote_password)
            # 创建 SFTP 会话
            sftp = ssh.open_sftp()
            # 上传文件
            sftp.put(self.image_save_path, self.remote_path + f"/{prompt}.jpg")
            print("文件上传成功")

        except Exception as e:
            # 如果失败，抛出异常
            raise ValueError(f"上传文件失败: {e}")
        finally:
            # 关闭 SFTP 会话和 SSH 连接。
            sftp.close()
            ssh.close()

    def post_request(
                    self,
                    init_image: np.array = [],
                    final_image: np.array = None,
                    prompt: str = ""
                    ):
        """发送请求计算抓取位置。
        
        Args:
            init_image_list (list[list[int]]): 初始图片的像素值列表.
            prompt (str): 用户发布给机器人的指令.
        
        Returns:
            action (list[int]): 抓取位置.
        """


        # if init_image == []:
        #     raise ValueError("init_image 不能为空")
        # if prompt == "":
        #     raise ValueError("prompt 不能为空")
        try:
            # 将图片保存至本地并上传到远程服务器。
            self.save_image_to_local(init_image, prompt)
            self.upload_image_to_remote(prompt)

            if final_image is not None:
                self.save_image_to_local(final_image, prompt + "_final")
                self.upload_image_to_remote(prompt + "_final")

            # 请求头。
            headers = {
                'Content-Type': 'application/json'
                }
            messages = [
                {
                    "role": "user",
                    "content":[

                        {
                            "type": "text",
                            "text":f"Based on the image, what action should the robot take to {prompt}"
                        },
                        {
                            "type": "image",
                            "image": self.remote_path + f"/{prompt}.jpg"
                        },
                    ]
                }
                ]

            
            # 请求体中需要使用远程服务器中的图片保存路径，远程服务器中图片保存的路径为 self.remote_path + f"/{prompt}.jpg"。
            payload = json.dumps({
                "messages": messages,
                "top_p": 0.7,
            })
            # 发送请求。
            # try:
            response = requests.request("POST", self.url, data=payload, headers=headers)
            response = response.text
            
            action = self.post_process(response)
            return action
        except Exception as e:
            raise ValueError(f"请求失败:{e}")


    def extract_action_list(self, response: str):
        import re
        import numpy as np
        # 使用正则表达式提取所有数字
        float_numbers = re.findall(r'-?\d+\.\d+', response)
        # 将提取到的数字转换为浮点数
        float_numbers = [float(num) for num in float_numbers]
        # 转为 action
        action = {"pose0_position": np.array(float_numbers[:2], dtype=np.float32), "pose1_position": np.array(float_numbers[2:], dtype=np.float32)}
        # 输出结果
        return action


    def post_process(self, response: str):
        """ 处理返回的抓取位置。
            返回值样例：
            这个指令的目的是将目标物品 pentagon 放置到容器 frame 中，所以我们需要分别获取目标物品的位置和容器的位置。
            步骤：
            步骤一：获取目标物品的位置，目标物品为 pentagon， 根据提供的场景图片，我们可以推算出 pentagon 的位置在 [0.5218750238418579, 0.25]。
            步骤二：获取目标容器的位置，目标物品为 frame， 根据提供的场景图片，我们可以推算出 frame 的位置在 [0.48499953746795654, -0.2838943600654602]。
            因此，机器人的动作序列为：
            {"pose0_position": [0.5218750238418579, 0.25], "pose1_position": [0.48499953746795654, -0.2838943600654602]}
            
        Args:
            response (str): 请求返回的内容.

        Returns:
            action (Dict[str, np.array]): 抓取位置和角度.
        """
        print(response)
        action = re.findall(r"<answer>\\n(.*?)\\n</answer>", response) 
        action = eval(action[0])
        dragged_obj = action["dragged_obj"]
        base_obj = action["base_obj"]
        actions = []
        for i in range(len(dragged_obj)):
            actions.append({"pose0_position": np.array(dragged_obj[i], dtype=np.float32), "pose1_position": np.array(base_obj[0], dtype=np.float32), "pose0_rotation": np.array([0.,0.,0.,0.9999], dtype=np.float32), "pose1_rotation": np.array([0.,0.,0.,0.9999], dtype=np.float32)} )
        
        return actions

    

# if __name__ == "__main__":
#     post_api = PostAPI(
#         url = "http://127.0.0.1:8080/qwen2_vl_service/chat",
#         remote_ip="connect.bjc1.seetacloud.com",
#         remote_user_name="root",
#         port=32219,
#         remote_password="mfH6cr4mtaPe",
#         remote_path="/root/autodl-tmp/qwen2_vl_service/images",
#         local_path="/home/zp/文档/code/robovlm/local_images"
#     )
